import numpy as np
import matplotlib.pyplot as plt


class LQRController:
    def __init__(self, A, B, Q, R):
        self.A = A
        self.B = B
        self.Q = Q
        self.R = R
        self.measurement_history = []
        self.setpoint_history = []

    def lqr_solve(self):
        """
        Solve the continuous time LQR controller.

        dx/dt = A*x + B*u

        cost = integral (x.T*Q*x + u.T*R*u) dt

        Parameters:
        A: System dynamics matrix
        B: Control input matrix
        Q: State cost matrix
        R: Control cost matrix

        Returns:
        K: State feedback gains
        P: Solution to Riccati equation
        eigVals: Eigenvalues of the closed loop system
        """
        # Check if the matrices are of compatible dimensions
        n = self.A.shape[0]
        m = self.B.shape[1]

        # Solve the continuous-time algebraic Riccati equation using a Lyapunov approach
        Z = np.vstack([
            np.hstack([self.A, -self.B @ np.linalg.inv(self.R) @ self.B.T]),
            np.hstack([-self.Q, -self.A.T])
        ])
        eigvals, eigvecs = np.linalg.eig(Z)

        # Select stable eigenvalues
        stable_indices = np.where(np.real(eigvals) < 0)[0]
        eigvecs_stable = eigvecs[:, stable_indices]

        # Separate eigvecs_stable into X and Y components
        X = eigvecs_stable[:n, :]
        Y = eigvecs_stable[n:, :]

        # Solve for P
        P = np.real(Y @ np.linalg.inv(X))

        # Compute the LQR gain
        K = np.linalg.inv(self.R) @ self.B.T @ P

        # Compute the closed loop eigenvalues
        eigVals, _ = np.linalg.eig(self.A - self.B @ K)

        return K, P, eigVals

    def update(self, setpoint, measurement):
        self.measurement_history.append(measurement)
        self.setpoint_history.append(setpoint)
        K, _, _ = self.lqr_solve()
        u = -K @ (measurement - setpoint)
        return u

    def plot_error(self, title='none'):
        self.setpoint_history = np.array(self.setpoint_history)
        self.measurement_history = np.array(self.measurement_history)
        for i in range(len(self.setpoint_history[0])):
            fig, ax = plt.subplots()
            ax.plot([i / 100 for i in range(len(self.setpoint_history[:,i]))], self.setpoint_history[:,i], label='setpoint')
            ax.plot([i / 100 for i in range(len(self.measurement_history[:,i]))], self.measurement_history[:,i], label='measurement')
            ax.set_xlabel('time')
            ax.set_ylabel('value')
            ax.set_title(title)
            ax.legend()


if __name__ == '__main__':
    A = np.array([[0., 1.], [-1., -1.]])
    B = np.array([[0.], [1.]])
    Q = np.array([[1., 0.], [0., 1.]])
    R = np.array([[1.]])

    controller = LQRController(A, B, Q, R)
    K, P, eigVals = controller.lqr_solve()
    u = controller.update(np.array([1, 0]), np.array([0, 0]))

    print("LQR gain matrix K:")
    print(K)
    print("\nSolution to Riccati equation P:")
    print(P)
    print("\nClosed loop system eigenvalues:")
    print(eigVals)
    print("Control value:")
    print(u)
